"""
    策略1， 根据 日线的close 的 fib 历史数据， 预测 1， 2 ，3 ，5  日的收益
"""
import pandas as pd
import numpy as np
import talib as ta
import datetime

from fib_config import stime, stime1, etime, etime1
from fib_utils import fib_yield_for
stock_list = []
all_flows = None


def fib_rate(fib_df, column, n=9):
    for i in fib_yield_for(n):
        fib_df[column + "_fib" + str(i)] = fib_df[column] / fib_df[column].shift(i) - 1.0


data = get_price(stock_list, stime, etime,
                         '1d',
                         #   close      涨跌幅        振幅         换手率
                         ['close', 'quote_rate', 'amp_rate', 'turnover_rate'],
                         True,  # 是否跳过停牌
                         "pre",  # 前复权
                         0,  # 天数
                         is_panel=1)
df = data.to_frame().reset_index()
# df.columns = ['date', 'symbol', 'close', 'quote_rate', 'amp_rate', 'turnover_rate']
df.rename(columns={'major': 'date','minor': 'symbol'}, inplace=True)
df = df.sort_values(['symbol', 'date'])
df = df.reset_index(drop=True)

# 因子
# 周线换手率均值, 月线换手率均值
df['close_ma5'] = df.groupby(['symbol'])['close'].rolling(5).mean().reset_index(drop=True, level=0)
df['close_ma20'] = df.groupby(['symbol'])['close'].rolling(20).mean().reset_index(drop=True, level=0)
# n天的平均价到今天的涨幅
df["close" + "_fib_d1"] = df["close"] / df["close"].shift(1) - 1.0
df["close" + "_fib_d2"] = df["close"] / df["close"].shift(2) - 1.0
df["close" + "_fib_d3"] = df["close"] / df["close"].shift(3) - 1.0
# n周的平均价到今天的涨幅
df["close" + "_fib_w1"] = df["close"] / df["close_ma5"].shift(1) - 1.0
df["close" + "_fib_w2"] = df["close"] / df["close_ma5"].shift(2) - 1.0
df["close" + "_fib_w3"] = df["close"] / df["close_ma5"].shift(3) - 1.0
# n月的平均价到今天的涨幅
df["close" + "_fib_m1"] = df["close"] / df["close_ma20"].shift(1) - 1.0
df["close" + "_fib_m2"] = df["close"] / df["close_ma20"].shift(2) - 1.0
df["close" + "_fib_m3"] = df["close"] / df["close_ma20"].shift(3) - 1.0

# 周线换手率均值, 月线换手率均值
df['turnover_r5'] = df.groupby(['symbol'])['turnover_rate'].rolling(5).mean().reset_index(drop=True, level=0)
df['turnover_r20'] = df.groupby(['symbol'])['turnover_rate'].rolling(20).mean().reset_index(drop=True, level=0)

df["turnover" + "_fib_d1"] = df.groupby(['symbol'])['turnover_rate'].shift(1).reset_index(drop=True, level=0)
df["turnover" + "_fib_d2"] = df.groupby(['symbol'])['turnover_rate'].shift(2).reset_index(drop=True, level=0)
df["turnover" + "_fib_d3"] = df.groupby(['symbol'])['turnover_rate'].shift(3).reset_index(drop=True, level=0)

df["turnover" + "_fib_w1"] = df.groupby(['symbol'])['turnover_r5'].shift(1).reset_index(drop=True, level=0)
df["turnover" + "_fib_w2"] = df.groupby(['symbol'])['turnover_r5'].shift(2).reset_index(drop=True, level=0)
df["turnover" + "_fib_w3"] = df.groupby(['symbol'])['turnover_r5'].shift(3).reset_index(drop=True, level=0)

df["turnover" + "_fib_m1"] = df.groupby(['symbol'])['turnover_r20'].shift(1).reset_index(drop=True, level=0)
df["turnover" + "_fib_m2"] = df.groupby(['symbol'])['turnover_r20'].shift(2).reset_index(drop=True, level=0)
df["turnover" + "_fib_m3"] = df.groupby(['symbol'])['turnover_r20'].shift(3).reset_index(drop=True, level=0)


# 收盘价比例， 结果中为NaN的地方表示空值
df['return1'] = df['close'] / df.groupby(['symbol'])['close'].shift(-1) - 1.0
df['return2'] = df['close'] / df.groupby(['symbol'])['close'].shift(-2) - 1.0
df['return3'] = df['close'] / df.groupby(['symbol'])['close'].shift(-3) - 1.0
df['return5'] = df['close'] / df.groupby(['symbol'])['close'].shift(-5) - 1.0
### 对 return 分组
# 删除NAN 的行
df = df.dropna(axis=0, how='any')
# 保留小数
round_map = {
    "quote_rate": 3, "amp_rate": 3, "turnover_rate": 3,
    "close_ma5": 3, "close_ma20": 3, "turnover_r5": 3, "turnover_r20": 3,

    "close_fib_d1": 3, "close_fib_d2": 3, "close_fib_d3": 3,
    "close_fib_w1": 3, "close_fib_w2": 3, "close_fib_w3": 3,
    "close_fib_m1": 3, "close_fib_m2": 3, "close_fib_m3": 3,

    "turnover_fib_d1": 3, "turnover_fib_d2": 3, "turnover_fib_d3": 3,
    "turnover_fib_w1": 3, "turnover_fib_w2": 3, "turnover_fib_w3": 3,
    "turnover_fib_m1": 3, "turnover_fib_m2": 3, "turnover_fib_m3": 3,

    "return1": 3, "return2": 3, "return3": 3, "return5": 3
}
# df = df.round(round_map)
df = df.round(3)

# 删除不需要的列, close 不是比例值
df.drop(['close', 'date', 'symbol',
         # "close_ma5", "close_ma20", "turnover_r5", "turnover_r20"
         ],
        axis=1,
        inplace=True)

df.to_csv(r'./fib_close_v1.csv', index=False)

